Description

Implements grouped-tensor quantize for the FP8 1D (1x128) and 2D (128x128) block-scaling recipes in row-wise (RW), column-wise (CW) and BOTH quantization directions. A single CUDA kernel launch walks 128x128 tiles across every tensor in the group, with each CTA decoding its owning tensor from the device-side GroupedTensor metadata with (N, R, K) shapes. Supports SAME_BOTH_DIMS (all tensors identical) and VARYING_FIRST_DIM (constant K, varying R) shape representations.

Three kernels share the dispatcher in group_quantize_blockwise_{1d,2d}:

• group_block_scaled_1d_rw_kernel — RW-only dispatch; 8 threads/row, reads global memory directly into vec-16 registers; bypasses TMA because the shared memory roundtrip and ptx::mbarrier does not buy anything without re-use in CW path.
• group_block_scaled_1d_tma_kernel — CW-only and BOTH dispatch; TMA bulk-load fills shared memory input cache. BOTH runs RW pass first (8 threads/row, vec-16 read from shared memory) then CW pass (2 threads/column, 64-row register stage); CW-only skips the RW pass. CW path writes the transposed-FP8 tile to a shared memory transpose staging buffer, then drains to global memory.
• group_block_scaled_2d_tma_kernel — RW-only, CW-only and BOTH dispatch; TMA bulk-load fills shared memory input cache. Pass 1 stages 8 IVecs/thread in registers while computing the per-tile scalar amax. Pass 2 quantizes from registers, emits row-wise output, stages column-wise output to shared memory transpose staging buffer, then drains to global memory.

Kernels are gated to Hopper (sm_90) at the host dispatcher (cuBlasLt grouped GEMM supports FP8 block-scaling only on Hopper).

PR includes PyTorch integration into te.GroupedTensor only.

PyTorch integration into te.GroupedLinear and JAX integration deferred to a follow-up PRs.

Partially resolves #2525

Performance

Benchmark on H200 with a sweep of grouped tensors in (N, M, K) shapes:

• N ∈ {4, 8, 16, 32, 64, 128} (# of device-local experts)
• M = 4096 @ N = 4 → M = 128 @ N = 128 (# of tokens/expert, scaling inversely with # experts)
• K ∈ {1024, 1792, 2048, 3584, 4096, 7168} (device-local shard of TP-hidden/intermediate-FFN dim)

Two shape families per config:

• U_MoE (uniform, SAME_BOTH_DIMS): all experts share the (M, K) shape
• J_MoE (jagged, VARYING_FIRST_DIM): per-expert M drawn from an imbalanced routing, common K

Buckets:

• Small/Unsaturated (S): R·K ≤ 32M elements (< 2048 tiles and < 15 waves on H200's 132 SMs)
• Large/Saturated (L): R·K > 32M elements (> 2048 tiles, SMs busy across many waves)

Bucket medians across 3 reps. Speedup is grouped vs the split-quantized fallback that loops over the grouped tensor and quantizes each constituent sequentially. % mono is grouped throughput relative to a single non-grouped FP8 block-scaling quantize on the equivalent monolithic (N·M, K) tensor.

Notes

• % of mono throughput is roughly consistent across buckets for every path, confirming no per-expert overhead in the new kernels.
• Greater-than-100% mono throughput cases come from TMA bulk-loads, register staging, and vec-16 reads that the non-grouped FP8 block-scaling kernels do not have.
• Speedup over split-quantize scales as expected with # of experts (roughly linearly in the unsaturated regime).
• S-bucket % memcpy is lower than L because launch and per-CTA setup are not amortized over a long bandwidth-bound steady state; absolute kernel times are still small (< 35 µs).

Known Sub-Optimalities

1D CW load bank conflicts on ~35% of load wavefronts (reading from the shared memory input-cache)

• No possible stride padding or XOR swizzle to alleviate.
• TMA hardware swizzle with CU_TENSOR_MAP_SWIZZLE_128B has the right pattern but caps FP16/BF16 at 64-elements; does not fit the 128-element tile for FP8 block-scaling without doubling per-tile launch overhead (quadrupling for FP32).
• CW-only stays at 2 t/col. Going to 4 t/col would double the bank-conflict footprint (4 lanes per column at the same row stride instead of 2), and CW-only is not occupancy-bound (3 CTAs/SM regardless), so the restructure costs more than it saves.
• 1D BOTH uses 4 t/col with a 32-row reg_data per thread and two column passes per CTA. The RW pass's per-expert scale-offset arithmetic plus a 64-row reg_data crossed the 85-reg / 3-CTA threshold on sm_90; halving reg_data restores 4 CTAs/SM. The doubled column pass and extra XOR-reduce stage are cheap relative to the occupancy gain.

1D BOTH reads the shared memory input-cache twice

• The RW (8 threads/row) and CW (2 threads/column) passes have different threading.
• Attempted to unify with 8 threads/row for both RW and CW. Caused bank conflicts on ~76% of store wavefronts (writing to the shared memory transpose buffer), reduced to ~43% with a XOR swizzle but not enough to beat separate RW/CW passes.
• Did not pursue the 2 threads/column unification; costs 40x more shfl ops than 8 threads/row attempt, plus a shared memory partial buffer and sync.

2D CW/BOTH has bank conflicts on ~16% of store wavefronts (when writing to the shared memory transpose buffer)

• Already reduced from ~75% via a XOR swizzle, further reduction was not possible.
• Minimal impact (< 5%) on kernel time.

No TMA-store

• MXFP8 grouped quantize kernel leverages this by decomposing a 128x128 tile into 32-row sub-stages that each have their own independent 32x1 or 1x32 scale; shared memory footprint is based on a single sub-stage; can be quantized and TMA-stored independently; hides TMA-store of one stage under the compute of next stage.
• FP8 block-scaling 128-element scale-block spans the entire 128-row tile. Cannot decompose into independent sub-stages and pipeline the TMA-stores. Single non-pipelined TMA-store requires holding the transposed staging buffer for the entire tile until all work on tile is finished, blows up shared memory footprint, collapses occupancy to 2CTA/SM. The recipe itself is the roadblock.

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes